import numpy as np
import math
from SE3 import Pose, matrix2quat

class Robot:
    def __init__(self, dh_model_list, dh_base_0, dh_end_tcp):  # DH:{alpha, a, theta, d}
        self.dof = len(dh_model_list)
        self.dh_0_1 = dh_model_list[0]
        self.dh_1_2 = dh_model_list[1]
        self.dh_2_3 = dh_model_list[2]
        self.dh_3_4 = dh_model_list[3]
        self.dh_4_5 = dh_model_list[4]
        self.dh_5_6 = dh_model_list[5]
        self.set_base(dh_base_0)
        self.set_tcp(dh_end_tcp)
        self.dh_a_list = [self.dh_0_1.a, self.dh_1_2.a, self.dh_2_3.a, self.dh_3_4.a, self.dh_4_5.a, self.dh_5_6.a]
        self.dh_d_list = [self.dh_0_1.d, self.dh_1_2.d, self.dh_2_3.d, self.dh_3_4.d, self.dh_4_5.d, self.dh_5_6.d]
    
    def set_base(self, dh):
        self.dh_base_0 = dh

    def set_tcp(self, dh):
        self.dh_end_tcp = dh

    def get_robot_tf_tree(self, joints):
        robot_tf_tree = [None, None, None, None, None, None]
        robot_tf_tree[0] = self.dh_0_1.to_matrix(joints[0])
        robot_tf_tree[1] = self.dh_1_2.to_matrix(joints[1])
        robot_tf_tree[2] = self.dh_2_3.to_matrix(joints[2])
        robot_tf_tree[3] = self.dh_3_4.to_matrix(joints[3])
        robot_tf_tree[4] = self.dh_4_5.to_matrix(joints[4])
        robot_tf_tree[5] = self.dh_5_6.to_matrix(joints[5])
        return robot_tf_tree
    
    def get_base_robot_tf(self):
        return self.dh_base_0.to_matrix(0)
        
    def get_end_tcp_tf(self):
        return self.dh_end_tcp.to_matrix(0)
    
    def get_dh_a_list(self):
        return self.dh_a_list
    
    def get_dh_d_list(self):
        return self.dh_d_list
    
    def set_robot_joint_limit(self, joints_limit):
        self.joint1_limit_min, self.joint1_limit_max = joints_limit[0]
        self.joint2_limit_min, self.joint2_limit_max = joints_limit[1]
        self.joint3_limit_min, self.joint3_limit_max = joints_limit[2]
        self.joint4_limit_min, self.joint4_limit_max = joints_limit[3]
        self.joint5_limit_min, self.joint5_limit_max = joints_limit[4]
        self.joint6_limit_min, self.joint6_limit_max = joints_limit[5]
        

if __name__ == '__main__':
    pass
